# IMPORT LIBRERIE
from ultralytics import YOLO
import os
import cv2
import torch
import numpy as np
from sklearn.metrics import precision_score, recall_score, f1_score, accuracy_score

# CONFIG
MODEL_PATH = "C:/Users/caruso/Desktop/YOLO_Project/agrilus/yolov11x/weights/best.pt"
TEST_IMAGES_PATH = "C:/Users/caruso/Desktop/YOLO_Project/images/test_random"  # Test image folder
GROUND_TRUTH_PATH = "C:/Users/caruso/Desktop/YOLO_Project/test_labels.txt"  # Files with real labels
RESULTS_FILE = "C:/Users/caruso/Desktop/YOLO_Project/risultati_valutazione.txt"  # File to save the results

# YOLO Class Mapping (names → numbers)
CLASS_MAPPING = {
    'Agrilus_angustulus': 0,
    'Agrilus_anxius': 1,
    'Agrilus_betuleti': 2,
    'Agrilus_cuprescens': 3,
    'Agrilus_graminis': 4,
    'Agrilus_hastulifer': 5,
    'Agrilus_laticornis': 6,
    'Agrilus_obscuricollis': 7,
    'Agrilus_olivicolor': 8,
    'Agrilus_planipennis': 9,
    'Agrilus_pratensis': 10,
    'Agrilus_sulcicollis': 11,
    'Agrilus_viridis': 12
}

CONF_THRESHOLD = 0.3  # Reduced to get more predictions

# 1. UPLOAD YOLO TEMPLATE
print("YOLO Template Upload...")
model = YOLO(MODEL_PATH)

# 2. UPLOAD TEST IMAGES
print("Loading test images...")
image_files = [f for f in os.listdir(TEST_IMAGES_PATH) if f.endswith(('.jpg', '.png'))]

# 3. UPLOAD REAL LABELS
print("Loading real labels...")
true_labels = {}
with open(GROUND_TRUTH_PATH, "r") as f:
    for line in f:
        image_name, class_id = line.strip().split()
        true_labels[image_name] = int(class_id)  # Convert to integer

# 4. INITIALIZE LISTS FOR METRICS
y_true = []  # Real lables
y_pred = []  # Model Predictions

# 5. MAKE PREDICTIONS ON ALL IMAGES
print("Making predictions on test images...")
for image_file in image_files:
    image_path = os.path.join(TEST_IMAGES_PATH, image_file)
    image = cv2.imread(image_path)

    # Run the prediction with YOLO
    results = model.predict(image, conf=CONF_THRESHOLD)

    # Check if there are valid predictions
    if results and len(results) > 0 and hasattr(results[0], "probs") and results[0].probs is not None:
        pred_index = torch.argmax(results[0].probs.data).item()
        pred_class_name = results[0].names[pred_index]
        pred_class = CLASS_MAPPING.get(pred_class_name, -1)  # Convert name to number
    else:
        pred_class = -1

    # Saving results
    if image_file in true_labels:
        y_true.append(true_labels[image_file])  # Real lable (number)
        y_pred.append(pred_class)  # Prediction converted into number

# 6. DEBUG: Prediction Check
print("Debug: Prediction Check")
print(f"Number of images tested: {len(y_true)}")
print(f"Valid predictions: {len(y_pred)}")

if not y_true or not y_pred:
    print("Error: No valid predictions.")
    exit()

# 7. CALCULATE METRICS
print("Calculating evaluation metrics...")
precision = precision_score(y_true, y_pred, average='weighted', zero_division=0)
recall = recall_score(y_true, y_pred, average='weighted', zero_division=0)
f1 = f1_score(y_true, y_pred, average='weighted', zero_division=0)
accuracy = accuracy_score(y_true, y_pred)

# 8. PRINT RESULTS
print("\n**Performance of YOLO**")
print(f"Precision: {precision:.4f}")
print(f"Recall: {recall:.4f}")
print(f"F1-score: {f1:.4f}")
print(f"Accuracy: {accuracy:.4f}")

# 9. SAVE RESULTS
with open(RESULTS_FILE, "w") as f:
    f.write("**Performances of YOLO**\n")
    f.write(f"Precision: {precision:.4f}\n")
    f.write(f"Recall: {recall:.4f}\n")
    f.write(f"F1-score: {f1:.4f}\n")
    f.write(f"Accuracy: {accuracy:.4f}\n")

print(f"Risults saved in: {RESULTS_FILE}")

